"""
在现有双阶段 YOLO 管道基础上，为数字增加 cnOCR(onnx, GPU) 兜底：

规则：
- 第一阶段与 ROI 构建完全复用 `DoubleYoloPipeline` 的逻辑；
- 第二阶段 YOLO 未能识别出 one~six 时，对 `number` ROI 使用 cnOCR 的单行识别模型
  `number-densenet_lite_136-fc` 进行兜底；
- cnOCR 识别出的字符若包含 1..6 中的一个数字，则映射为 one~six 并作为最终数字结果；
- 其他流程（cone 颜色融合、dashboard 状态）保持不变。
"""

import os
import importlib.util
import importlib.machinery
from typing import Any, Dict, List, Optional, Tuple, Union

import cv2
import numpy as np

def _load_base_pipeline_class() -> Any:
    """动态加载同目录下的 `double_yolo_pipeline-1.py`，返回 DoubleYoloPipeline 类。
    之所以动态加载，是因为文件名包含连字符，无法用常规 import。"""
    base_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(base_dir, "double_yolo_pipeline-1.py")
    module_name = "double_yolo_pipeline_1"
    loader = importlib.machinery.SourceFileLoader(module_name, file_path)
    spec = importlib.util.spec_from_loader(module_name, loader)
    if spec is None:
        raise ImportError("Cannot create spec for double_yolo_pipeline-1.py")
    module = importlib.util.module_from_spec(spec)
    loader.exec_module(module)  # type: ignore[attr-defined]
    if not hasattr(module, "DoubleYoloPipeline"):
        raise ImportError("DoubleYoloPipeline not found in double_yolo_pipeline-1.py")
    return getattr(module, "DoubleYoloPipeline")


BaseDoubleYoloPipeline = _load_base_pipeline_class()


class DoubleYoloPipelineWithCnOCR(BaseDoubleYoloPipeline):
    def __init__(self,
                 model_path: str,
                 device: Optional[str] = None,
                 debug: bool = True,
                 save_dir: Optional[str] = None,
                 enable_cnocr: bool = True) -> None:
        super().__init__(model_path=model_path, device=device, debug=debug, save_dir=save_dir)
        self.enable_cnocr = enable_cnocr
        self._cnocr: Optional[Any] = None

    # ========================= cnOCR 相关 =========================
    def _ensure_cnocr(self) -> None:
        if not self.enable_cnocr or self._cnocr is not None:
            return
        try:
            from cnocr import CnOcr  # type: ignore
        except Exception as e:
            self._debug_print(f"⚠️ 未安装 cnocr，跳过兜底: {e}")
            self.enable_cnocr = False
            return

        # 兼容不同版本的参数命名：优先尝试 onnx + GPU
        tried: List[str] = []
        exc: Optional[Exception] = None
        for kwargs in [
            {"rec_model_name": "number-densenet_lite_136-fc", "use_gpu": True, "model_backend": "onnx"},
            {"rec_model_name": "number-densenet_lite_136-fc", "use_gpu": True, "onnx": True},
            {"rec_model_name": "number-densenet_lite_136-fc", "use_gpu": True},
            {"rec_model_name": "number-densenet_lite_136-fc", "use_gpu": False},
        ]:
            try:
                self._cnocr = CnOcr(**kwargs)
                self._debug_print(f"✅ cnOCR 初始化成功: {kwargs}")
                exc = None
                break
            except Exception as e:  # pragma: no cover - 容错初始化
                tried.append(str(kwargs))
                exc = e
                self._cnocr = None
                continue
        if self._cnocr is None:
            self._debug_print(f"⚠️ cnOCR 初始化失败，已尝试: {tried}; 错误: {exc}")
            self.enable_cnocr = False

    @staticmethod
    def _map_digit_to_word(text: str) -> Optional[str]:
        if not text:
            return None
        mapping = {
            "1": "one",
            "2": "two",
            "3": "three",
            "4": "four",
            "5": "five",
            "6": "six",
        }
        # 提取第一个 1..6 的数字字符
        for ch in text:
            if ch in mapping:
                return mapping[ch]
        return None

    def _cnocr_single_line_digits(self, roi_bgr: np.ndarray) -> Tuple[str, float]:
        """对数字 ROI 做单行 OCR，返回 (name, score)。失败时返回 ("None", 0.0)。"""
        if roi_bgr is None or roi_bgr.size == 0:
            return ("None", 0.0)
        self._ensure_cnocr()
        if not self.enable_cnocr or self._cnocr is None:
            return ("None", 0.0)

        try:
            # cnOCR 接受 BGR/RGB/灰度 numpy 数组；保持尺寸即可
            roi = roi_bgr
            out = self._cnocr.ocr_for_single_line(roi)  # 兼容不同版本，返回可能是 list 或 dict
            text: str = ""
            score: float = 0.0
            if isinstance(out, list):
                # 新版通常返回 list[ {"text": str, "score": float} ]
                if len(out) > 0:
                    item = out[0]
                    if isinstance(item, dict):
                        text = str(item.get("text", ""))
                        score = float(item.get("score", 0.0))
                    else:
                        text = str(item)
                        score = 0.0
            elif isinstance(out, dict):
                text = str(out.get("text", ""))
                score = float(out.get("score", 0.0))
            else:
                text = str(out)
                score = 0.0

            mapped = self._map_digit_to_word(text)
            if mapped is not None:
                # 若无置信度信息，给一个保守默认分数
                score_final = score if score and score > 0 else 0.85
                return (mapped, float(score_final))
            return ("None", 0.0)
        except Exception as e:  # pragma: no cover - 推理容错
            self._debug_print(f"⚠️ cnOCR 推理失败: {e}")
            return ("None", 0.0)

    # ========================= 覆写流程入口（仅加入兜底） =========================
    def process_image(self,
                      image: np.ndarray,
                      top_cut_ratio: float = 0.25,
                      width_center_ratio: float = 0.5,
                      first_stage_conf: float = 0.6,
                      second_stage_conf: float = 0.5,
                      final_conf: float = 0.6,
                      frame_idx: int = 0) -> Dict[str, Any]:
        if image is None or image.size == 0:
            raise ValueError("Input image is empty")
        try:
            cropped = self.crop_image(image, top_cut_ratio, width_center_ratio)
            first_stage_results = self.yolo_infer(cropped, conf=first_stage_conf, imgsz=640, debug_title_prefix="first-stage")
            detections = self.analyze_results(first_stage_results, conf=first_stage_conf)
            rois_with_boxes = self.build_adaptive_rois_with_boxes(detections, cropped)

            # 准备数字 ROI 给第二阶段
            number_roi = None
            for (cls_name, roi_img, _bbox) in rois_with_boxes:
                if cls_name == "number":
                    number_roi = roi_img
                    break
            second_stage_number_res: Tuple[str, float, Optional[List[float]]] = ("None", 0.0, None)
            if number_roi is not None:
                second_stage_number_res = self.second_stage_number(number_roi, conf=second_stage_conf)

            # ========= 新增：cnOCR 兜底 =========
            n_name, n_score, n_box = second_stage_number_res
            if n_name not in self.NUMBER_CLASSES and number_roi is not None:
                ocr_name, ocr_score = self._cnocr_single_line_digits(number_roi)
                if ocr_name in self.NUMBER_CLASSES:
                    # 用 OCR 结果替代
                    second_stage_number_res = (ocr_name, float(ocr_score), None)
                    self._debug_print(f"🧩 二阶段未命中，采用 cnOCR 结果: {ocr_name} ({ocr_score:.2f})")

            # HSV 对锥桶 ROI 投票（稳健版，含回退）
            hsv_vote_name: Optional[str] = None
            hsv_vote_conf: float = 0.0
            for (cls_name, roi_img, _bbox) in rois_with_boxes:
                if cls_name == "cone":
                    name_rb, score_rb, _ = self._vote_color_robust(roi_img)
                    if name_rb in {"red", "orange", "yellow", "blue"}:
                        hsv_vote_name, hsv_vote_conf = name_rb, float(score_rb)
                    else:
                        best_c, counts = self._count_hsv_by_ranges(roi_img)
                        total = max(1, sum(counts.values()))
                        hsv_vote_name = best_c if best_c in {"red", "orange", "yellow", "blue"} else None
                        hsv_vote_conf = float(counts.get(best_c, 0)) / float(total)
                    break

            if hsv_vote_name is None:
                for (cls_name, _roi_img, (bx1, by1, bx2, by2)) in rois_with_boxes:
                    if cls_name == "number":
                        h_n = max(1, by2 - by1)
                        cy2 = by1
                        cy1 = self._clip(by1 - h_n, 0, cropped.shape[0])
                        cx1, cx2 = bx1, bx2
                        if cy2 > cy1 and cx2 > cx1:
                            cone_roi_retry = cropped[cy1:cy2, cx1:cx2]
                            name_rb, score_rb, _ = self._vote_color_robust(cone_roi_retry)
                            if name_rb in {"red", "orange", "yellow", "blue"}:
                                hsv_vote_name, hsv_vote_conf = name_rb, float(score_rb)
                            else:
                                best_c, counts = self._count_hsv_by_ranges(cone_roi_retry)
                                total = max(1, sum(counts.values()))
                                hsv_vote_name = best_c if best_c in {"red", "orange", "yellow", "blue"} else None
                                hsv_vote_conf = float(counts.get(best_c, 0)) / float(total)
                        break

            final_results = self.finalize_results(
                first_stage_results,
                detections,
                second_stage_number_res,
                hsv_cone_vote=(hsv_vote_name, hsv_vote_conf) if hsv_vote_name is not None else None,
                conf=final_conf,
            )

            # 结构化结果
            cone_color = "None"
            number = "None"
            dashboard_state = "None"
            for class_name, detected_name, _confidence in final_results:
                if class_name == "cone":
                    cone_color = str(detected_name)
                elif class_name == "number":
                    number = str(detected_name)
                elif class_name == "dashboard":
                    dashboard_state = str(detected_name)

            # 可视化
            self.visualize_pipeline(frame_idx, image, cropped, detections, rois_with_boxes, second_stage_number_res, final_results)

            return {
                'final_results': final_results,
                'cone_color': cone_color,
                'number': number,
                'dashboard_state': dashboard_state,
                'processing_info': {
                    'image_shape': image.shape,
                    'cropped_shape': cropped.shape,
                    'num_detections_stage1': len(detections),
                    'num_rois': len(rois_with_boxes),
                    'roi_types': [cls_name for (cls_name, _roi_img, _bbox) in rois_with_boxes],
                    'crop_offsets': {
                        'top': int(self._last_crop_offsets.get('top', 0)),
                        'left': int(self._last_crop_offsets.get('left', 0)),
                    },
                    'cnocr_enabled': bool(self.enable_cnocr),
                }
            }
        except Exception as e:
            err = f"管道处理中出错: {e}"
            self._debug_print(err)
            return {
                'final_results': [],
                'cone_color': "None",
                'number': "None",
                'dashboard_state': "None",
                'error': err,
                'processing_info': {}
            }


def create_double_yolo_pipeline_with_cnocr(model_path: str,
                                           device: Optional[str] = None,
                                           debug: bool = True,
                                           save_dir: Optional[str] = None,
                                           enable_cnocr: bool = True):
    return DoubleYoloPipelineWithCnOCR(model_path=model_path,
                                       device=device,
                                       debug=debug,
                                       save_dir=save_dir,
                                       enable_cnocr=enable_cnocr)


